use std::fs::File;
use std::io::Write;

use cosmic_text::{SwashCache, FontSystem, Metrics, BorrowedWithFontSystem, Wrap};
use fontdb::{FaceInfo, Family};
use rand::rngs::ThreadRng;
use rand::Rng;
use tiny_skia::{Pixmap, ColorU8};
use png::Encoder;

const W: u32 = 300;
const H: u32 = 150;
const PT: f32 = 32.0;
const MIN_PT: f32 = 17.0;

fn main() {
    _run();
    // _test();
}
fn _run() {
    // images preparation
    let mut pixmap = Pixmap::new(W, H).unwrap();
    let mut rng = rand::thread_rng();
    let mut fonts = fontdb::Database::new();
    fonts.load_fonts_dir("files/fonts");
    let mut font_system = FontSystem::new_with_locale_and_db("en-US".to_owned(), fonts);
    let faces: Vec<FaceInfo> = font_system.db().faces()
        .map(|face| face.clone())
        .collect();
    let mut swash_cache = SwashCache::new();
    let mut text_buffer = cosmic_text::Buffer::new(&mut font_system, Metrics::new(PT, PT));
    text_buffer.set_size(&mut font_system, W as f32, H as f32);
    text_buffer.set_wrap(&mut font_system, Wrap::None);

    // text preparation
    let mut train_text = String::new();
    let mut test_text = String::new();

    // datasets generation
    // training
    println!("generating training set");
    for i in 0..512 {
        generate(
            &mut train_text,
            &faces, 
            &mut pixmap,
            &mut rng, 
            &mut swash_cache, 
            &mut text_buffer.borrow_with(&mut font_system)
        );
        save_png(&pixmap, &format!("files/train/{}.png", i));
    }
    let mut text_file = File::create("files/train/text.txt").unwrap();
    text_file.write_all(train_text.as_bytes()).unwrap();
    // testing
    println!("generating testing set");
    for i in 0..64 {
        generate(
            &mut test_text, 
            &faces, 
            &mut pixmap, 
            &mut rng, 
            &mut swash_cache, 
        &mut text_buffer.borrow_with(&mut font_system));
        save_png(&pixmap, &format!("files/test/{}.png", i));
    }
    let mut text_file = File::create("files/test/text.txt").unwrap();
    text_file.write_all(test_text.as_bytes()).unwrap();
}

fn _test() {
}

const MIN_L: usize = 27;
const MAX_L: usize = 31;

fn generate<'a>(
    texts: &mut String,
    faces: &[FaceInfo],
    pixmap: &mut Pixmap, 
    rng: &mut ThreadRng,
    swash_cache: &mut SwashCache,
    buffer: &mut BorrowedWithFontSystem<'a, cosmic_text::Buffer>
) {
    // random text
    let len = rng.gen_range(MIN_L..MAX_L);
    let text = (0..len)
        .map(|_i| rng.gen_range(48..58))
        .collect::<Vec<u8>>();
    let text = String::from_utf8(text).unwrap();
    texts.push_str("\n\u{0002}");
    texts.push_str(&text);
    texts.push('\u{0003}');

    // random font
    let attrs = face_to_attrs(&faces[rng.gen_range(0..faces.len())]);
    buffer.set_text(&text, attrs, cosmic_text::Shaping::Advanced);

    // random size and position;
    buffer.set_metrics(Metrics::new(PT, PT));
    let layout = &buffer.line_layout(0).unwrap()[0];
    let w = layout.w;
    let h = layout.max_ascent + layout.max_descent;

    let max_pt = f32::min(PT / w * W as f32, PT / h * H as f32);

    let pt = if max_pt > MIN_PT {
        rng.gen_range(MIN_PT..=max_pt)
    } else {
        max_pt
    };
    buffer.set_metrics(Metrics::new(pt, pt));
    let layout = &buffer.line_layout(0).unwrap()[0];
    let w = layout.w;
    let h = layout.max_ascent + layout.max_descent;

    let start_x = rng.gen_range(0.0..(W as f32 - w + 0.01)) as i32;
    let start_y = rng.gen_range(0.0..(H as f32 - h + 0.01)) as i32;

    buffer.shape_until_scroll();

    pixmap.fill(tiny_skia::Color::WHITE);
    let text_color = cosmic_text::Color::rgba(0, 0, 0, 0xFF);
    buffer.draw(swash_cache, text_color, |x, y, _w, _h, color| {
        draw_pixel(pixmap, color.as_rgba(), (x + start_x) as usize, (y + start_y) as usize)
    });
}

// convert a FaceInfo to one Attrs, ignoring other families
fn face_to_attrs<'a>(face: &'a FaceInfo) -> cosmic_text::Attrs<'a> {
    let family = &face.families[0].0;
    cosmic_text::Attrs {
        color_opt: None,
        family: Family::Name(family),
        stretch: face.stretch,
        style: face.style,
        weight: face.weight,
        metadata: 0
    }
}

// draw a pixel using alpha blending
fn draw_pixel(pixmap: &mut Pixmap, color: [u8; 4], x: usize, y: usize) {
    let color = tiny_skia::ColorU8::from_rgba(color[0], color[1], color[2], color[3]);
    let index = x + y * pixmap.width() as usize;
    let data = pixmap.pixels_mut();
    let dst = data[index].demultiply();
    let premultiplied_src = color.premultiply();
    let dst_alpha = 255 - color.alpha();
    let result = ColorU8::from_rgba(
        multiply_a(dst.red(), dst_alpha) + premultiplied_src.red(), 
        multiply_a(dst.green(), dst_alpha) + premultiplied_src.green(), 
        multiply_a(dst.blue(), dst_alpha) + premultiplied_src.blue(), 
        255
    );
    data[index] = result.premultiply();
}
fn multiply_a(c: u8, a: u8) -> u8 {
    let prod = u32::from(c) * u32::from(a) + 128;
    ((prod + (prod >> 8)) >> 8) as u8
}

// save RGB png file
fn save_png(image: &Pixmap, path: &str) {
    let file = File::create(path).unwrap();
    let ref mut w = std::io::BufWriter::new(file);
    let mut encoder = Encoder::new(w, image.width(), image.height());
    encoder.set_color(png::ColorType::Rgb);
    let mut writer = encoder.write_header().unwrap();

    let data: Vec<u8> = image.pixels()
        .iter()
        .map(|c| {
            let c = c.demultiply();
            [c.red(), c.green(), c.blue()]
        })
        .flatten()
        .collect();
    writer.write_image_data(&data).unwrap();
}